#include "checkers/LockGranularityChecker.h"

extern ReportSpace::Report external_bug_report;
unordered_map<FunctionDecl *, bool> reentrantMap;
struct code
{

    string fileName = "";
    unsigned line = 0;
    unsigned col = 0;
    string var = "";

    void dump()
    {
        cout << "code dump :: line : " << line << " ,";
    }
};
class LockInfo
{
public:
    unsigned beginLine;
    unsigned beginCol;
    unsigned endLine;
    unsigned endCol;

    vector<code> report;

    LockInfo()
    {
    }
};

class Reentrant : public ASTConsumer, public RecursiveASTVisitor<Reentrant>
{
public:
    ASTContext *context;
    SourceManager *sm;
    ASTManager *manager;
    ASTResource *resource;
    set<VarDecl *> globalVarSet;
    bool state = true;
    Reentrant(ASTContext *astcontext, ASTManager *astmanager, ASTResource *astresource, set<VarDecl *> varList) : context(astcontext), manager(astmanager), resource(astresource), globalVarSet(varList)
    {
        sm = &(context->getSourceManager());
    }
    void HandleTranslationUnit(ASTContext &cxt) override
    {
        context = &cxt;
        TranslationUnitDecl *decl = cxt.getTranslationUnitDecl();
        TraverseDecl(decl);
    }
    void findReentrant(FunctionDecl *F)
    {
        TraverseDecl(F);
    }
    //TODO
    bool VisitDeclRefExpr(DeclRefExpr *DRE)
    {

        if (DRE->getDecl())
        {
            if (VarDecl *VD = dyn_cast<VarDecl>(DRE->getDecl()))
            {
                if (globalVarSet.find(VD) != globalVarSet.end())
                {
                    state = false;
                    return false;
                }
            }
        }
        return true;
    }

    bool getState()
    {
        return state;
    }
};

class GlobalVarDecl : public ASTConsumer, public RecursiveASTVisitor<GlobalVarDecl>
{
public:
    set<VarDecl *> records;

    ASTContext *Context;
    set<VarDecl *> &getSet()
    {
        return records;
    }
    void HandleTranslationUnit(ASTContext &cxt) override
    {
        Context = &cxt;
        TranslationUnitDecl *decl = cxt.getTranslationUnitDecl();
        TraverseDecl(decl);
    }
    bool VisitVarDecl(VarDecl *VD)
    {
        clang::SourceManager &sm = Context->getSourceManager();
        if (!VD)
            return true;
        if (!sm.isInMainFile(VD->getLocation()) && std::string(VD->clang::Decl::getDeclKindName()) != "TranslationUnit")
        {
            return true;
        }
        string loc = VD->getBeginLoc().printToString(sm);
        if (loc.find("/lib/modules/") != string::npos || loc.find("/usr/src/") != string::npos)
        {
            return true;
        }
        if (const ParmVarDecl *pv = dyn_cast<ParmVarDecl>(VD))
        {
            return true;
        }
        if (VD->hasGlobalStorage() && !VD->isStaticLocal())
        {
            records.insert(VD);
        }
        return true;
    }
};

class NotRelevantStmtVisiter : public clang::ASTConsumer, public clang::RecursiveASTVisitor<NotRelevantStmtVisiter>
{
private:
    ASTContext *context;
    SourceManager *sm;
    ASTManager *manager;
    ASTResource *resource;
    FunctionDecl *FD;
    set<string> opList;
    vector<LockInfo *> *lockList;
    int level;
    set<string> lockFindList;
    set<string> unlockFindList;
    vector<LockInfo *> lockStack;
    set<VarDecl *> globalVarSet;

public:
    NotRelevantStmtVisiter(ASTContext *astcontext, ASTManager *astmanager, ASTResource *astresource, FunctionDecl *functiondecl, vector<LockInfo *> *lock, set<VarDecl *> varList, set<string> list, set<string> list2) : context(astcontext), manager(astmanager), resource(astresource), FD(functiondecl), lockList(lock), globalVarSet(varList), lockFindList(list), unlockFindList(list2)
    {
        sm = &(context->getSourceManager());
        level = 0;
        for (auto it = FD->param_begin(), end = FD->param_end(); it != end; it++)
        {
            auto iter = *it;
            globalVarSet.insert(iter);
        }
    }
    void HandleTranslationUnit(clang::ASTContext &Context)
    {

        //TraverseDecl(FD);
        if (!FD->hasBody())
        {
            return;
        }
        if (CompoundStmt *st = dyn_cast<CompoundStmt>(FD->getBody()))
        {
            if (st->body_empty())
            {
                return;
            }

            handleCompoundStmt(st);
        }
    }

    void handleCompoundStmt(CompoundStmt *s)
    {
        if (s->body_empty())
        {
            return;
        }
        for (auto iter = s->body_begin(), end = s->body_end(); iter != end; iter++)
        {
            Stmt *st = *iter;
            if (st == nullptr)
            {
                continue;
            }
            if (!handleStmt(st))
            {
                if (level != 0)
                {
                    code tmp;
                    fileLoc loc = getLocFromString(st->getBeginLoc().printToString(*sm));
                    tmp.fileName = loc.name;
                    tmp.line = loc.line;
                    tmp.col = loc.col;
                    lockStack.back()->report.push_back(tmp);
                }
            }
        }
        return;
    }

    bool handleStmt(Stmt *stmt)
    {

        if (stmt == nullptr)
        {
            return true;
        }

        if (CompoundStmt *s = dyn_cast<CompoundStmt>(stmt))
        {
            handleCompoundStmt(s);
            return true;
        }

        bool shared = false;

        shared = handleSubStmt(stmt);

        return shared;
    }

    bool handleSubStmt(Stmt *stmt)
    {
        bool shared = false;
        if (!stmt)
        {
            return shared;
        }

        if (BinaryOperator *BO = dyn_cast<BinaryOperator>(stmt))
        {
            handleBinaryOperator(BO);
        }

        if (DeclStmt *DS = dyn_cast<DeclStmt>(stmt))
        {
            handleDeclStmt(DS);
        }

        if (CallExpr *CE = dyn_cast<CallExpr>(stmt))
        {
            shared = handleCallExpr(CE);
            return shared;
        }

        if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(stmt))
        {
            shared = handleDeclRefExpr(DRE);
            return shared;
        }

        for (auto iter = stmt->child_begin(), end = stmt->child_end(); iter != end; iter++)
        {
            Stmt *st = *iter;
            if (st == nullptr)
            {
                continue;
            }
            if (CompoundStmt *c = dyn_cast<CompoundStmt>(st))
            {
                handleCompoundStmt(c);
                continue;
            }
            shared |= handleSubStmt(st);
        }

        return shared;
    }

    void handleBinaryOperator(BinaryOperator *BO)
    {
        if (!BO->isAssignmentOp())
        {
            return;
        }
        Expr *lhs = BO->getLHS();
        Expr *rhs = BO->getRHS();
        if (!lhs || !rhs)
        {
            return;
        }
        if (DeclRefExpr *dre = dyn_cast<DeclRefExpr>(lhs))
        {
            if (dre->getType().getTypePtr()->isPointerType())
            {
                queue<Stmt *> q;
                q.push(rhs);

                while (!q.empty())
                {
                    auto top = q.front();
                    q.pop();
                    if (!top)
                    {
                        continue;
                    }
                    if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(top))
                    {
                        if (DRE->getDecl())
                        {
                            if (VarDecl *VD = dyn_cast<VarDecl>(DRE->getDecl()))
                            {
                                if (globalVarSet.find(VD) != globalVarSet.end())
                                {
                                    if (dre->getDecl())
                                    {
                                        if (VarDecl *v = dyn_cast<VarDecl>(dre->getDecl()))
                                        {
                                            globalVarSet.insert(v);
                                            return;
                                        }
                                    }
                                }
                            }
                        }
                    }
                    auto it = top->child_begin();
                    for (; it != top->child_end(); ++it)
                    {
                        q.push(*it);
                    }
                }
            }
        }
    }

    void handleDeclStmt(DeclStmt *DS)
    {
        if (!DS->isSingleDecl())
        {
            return;
        }
        if (!DS->getSingleDecl())
        {
            return;
        }
        if (VarDecl *v = dyn_cast<VarDecl>(DS->getSingleDecl()))
        {
            if (v->getType().getTypePtr()->isPointerType())
            {

                queue<Stmt *> q;
                q.push(DS);

                while (!q.empty())
                {
                    auto top = q.front();
                    q.pop();
                    if (!top)
                    {
                        continue;
                    }
                    if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(top))
                    {
                        if (DRE->getDecl())
                        {
                            if (VarDecl *VD = dyn_cast<VarDecl>(DRE->getDecl()))
                            {

                                if (globalVarSet.find(VD) != globalVarSet.end())
                                {

                                    globalVarSet.insert(v);
                                    return;
                                }
                            }
                        }
                    }
                    auto it = top->child_begin();
                    for (; it != top->child_end(); ++it)
                    {
                        q.push(*it);
                    }
                }
            }
        }
    }

    // is FunctionDEcl Reentrant
    bool isReentrant(FunctionDecl *F)
    {
        bool state = true;
        Reentrant R(context, manager, resource, globalVarSet);
        R.findReentrant(F);
        return R.getState();
    }

    bool handleCallExpr(CallExpr *CE)
    {

        bool shared = false;
        if (!CE)
        {
            return shared;
        }

        FunctionDecl *callee = CE->getDirectCallee();
        bool state;
        if (reentrantMap.find(callee) != reentrantMap.end())
        {
            state = reentrantMap[callee];
        }
        else
        {
            state = isReentrant(callee);
            reentrantMap[callee] = state;
        }
        if (!state)
        {
            return true;
        }

        ArrayRef<Stmt *> AR = CE->getRawSubExprs();

        for (auto it = AR.begin(), end = AR.end(); it != end; it++)
        {
            Stmt *st = *it;

            queue<Stmt *> q;
            q.push(st);

            while (!q.empty())
            {
                auto top = q.front();
                q.pop();
                if (!top)
                {
                    continue;
                }

                if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(top))
                {

                    if (DRE->getDecl())
                    {
                        if (VarDecl *VD = dyn_cast<VarDecl>(DRE->getDecl()))
                        {

                            if (globalVarSet.find(VD) != globalVarSet.end())
                            {
                                shared = true;
                                return shared;
                            }
                        }
                    }
                }

                auto it = top->child_begin();
                for (; it != top->child_end(); ++it)
                {
                    q.push(*it);
                }
            }
            if (MemberExpr *ME = dyn_cast<MemberExpr>(st))
            {
                string name = ME->getMemberDecl()->getNameAsString();

                if (lockFindList.find(name) != lockFindList.end())
                {

                    LockInfo *tmp = new LockInfo();
                    tmp->beginLine = context->getFullLoc(ME->getBeginLoc()).getLineNumber();
                    tmp->beginCol = context->getFullLoc(ME->getBeginLoc()).getColumnNumber();
                    lockList->push_back(tmp);
                    lockStack.push_back(tmp);
                    level++;
                    shared = true;
                    continue;
                }
                if (unlockFindList.find(name) != unlockFindList.end())
                {

                    if (level == 0)
                    {
                        continue;
                    }
                    lockStack.back()->endLine = context->getFullLoc(ME->getEndLoc()).getLineNumber();
                    lockStack.back()->endCol = context->getFullLoc(ME->getEndLoc()).getColumnNumber();
                    lockStack.pop_back();
                    level--;
                    shared = true;
                    continue;
                }
                if (!ME->getBase()->IgnoreCasts())
                {
                    continue;
                }
                if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(ME->getBase()->IgnoreCasts()))
                {
                    if (VarDecl *VD = dyn_cast<VarDecl>(DRE->getDecl()))
                    {
                        if (globalVarSet.find(VD) != globalVarSet.end())
                        {
                            shared = true;
                            continue;
                        }
                    }
                }
            }
            else if (Expr *E = dyn_cast<Expr>(st))
            {
                if (!E->IgnoreCasts())
                {
                    continue;
                }
                if (DeclRefExpr *DRE = dyn_cast<DeclRefExpr>(E->IgnoreCasts()))
                {
                    if (!DRE->getDecl())
                    {
                        continue;
                    }
                    if (VarDecl *VD = dyn_cast<VarDecl>(DRE->getDecl()))
                    {
                        if (globalVarSet.find(VD) != globalVarSet.end())
                        {
                            shared = true;
                            continue;
                        }
                    }
                    else if (FunctionDecl *F = dyn_cast<FunctionDecl>(DRE->getDecl()))
                    {
                        string name = F->getNameAsString();
                        if (lockFindList.find(name) != lockFindList.end())
                        {

                            LockInfo *tmp = new LockInfo();
                            tmp->beginLine = context->getFullLoc(CE->getBeginLoc()).getLineNumber();
                            tmp->beginCol = context->getFullLoc(CE->getBeginLoc()).getColumnNumber();
                            lockList->push_back(tmp);
                            lockStack.push_back(tmp);
                            level++;
                            shared = true;
                            continue;
                        }
                        if (unlockFindList.find(name) != unlockFindList.end())
                        {

                            if (level == 0)
                            {
                                continue;
                            }
                            lockStack.back()->endLine = context->getFullLoc(CE->getEndLoc()).getLineNumber();
                            lockStack.back()->endCol = context->getFullLoc(CE->getEndLoc()).getColumnNumber();
                            lockStack.pop_back();
                            level--;
                            shared = true;
                            continue;
                        }
                    }
                }
            }
        }

        return shared;
    }

    bool handleDeclRefExpr(DeclRefExpr *DRE)
    {
        bool shared = false;
        if (!DRE)
        {
            return shared;
        }

        string kind = DRE->getDecl()->getCanonicalDecl()->getDeclKindName();
        if (kind.compare("Function") == 0)
        {
            return shared;
        }
        Decl *D = DRE->getDecl()->getCanonicalDecl();
        if (!D)
        {
            return shared;
        }
        VarDecl *VD = dyn_cast<VarDecl>(D);
        if (!VD)
        {
            return shared;
        }

        if (globalVarSet.find(VD) != globalVarSet.end() && VD->getType().getAsString().compare("std::mutex") != 0)
        {
            shared = true;
        }

        return shared;
    }

private:
    struct fileLoc
    {
        string name;
        unsigned line;
        unsigned col;
    };
    fileLoc getLocFromString(string fileName)
    {
        fileLoc loc;
        size_t pos = fileName.find(":");
        loc.name = fileName.substr(0, pos);
        string tmp = fileName.substr(pos + 1);
        pos = tmp.find(":");
        string lineStr = tmp.substr(0, pos);
        string tmp2 = tmp.substr(pos + 1);
        pos = tmp2.find(" ");
        string colStr = tmp2.substr(0, pos);
        loc.line = (unsigned)stoi(lineStr);
        loc.col = (unsigned)stoi(colStr);
        return loc;
    }
};

class LockVisiter : public clang::ASTConsumer, public clang::RecursiveASTVisitor<LockVisiter>
{
private:
    ASTContext *context;
    SourceManager *sm;
    ASTManager *manager;
    ASTResource *resource;

    vector<LockInfo *> lockList;
    set<string> banList = {"func"};
    set<string> lockFindList = {"lock"};
    set<string> unlockFindList = {"unlock"};
    set<VarDecl *> globalVarDeclSet;

public:
    LockVisiter(ASTContext *astcontext, ASTManager *astmanager, ASTResource *astresource, set<VarDecl *> vardeclset, set<string> lock, set<string> unlock) : context(astcontext), manager(astmanager), resource(astresource), globalVarDeclSet(vardeclset), lockFindList(lock), unlockFindList(unlock)
    {
        sm = &(context->getSourceManager());
    }

    void HandleTranslationUnit(clang::ASTContext &Context)
    {
        TraverseDecl(Context.getTranslationUnitDecl());
    }

    vector<LockInfo *> &getList()
    {
        return lockList;
    }
    bool VisitFunctionDecl(FunctionDecl *FD)
    {
        if (!FD->hasBody())
        {
            return true;
        }
        if (!sm->isInMainFile(FD->getLocation()) && std::string(FD->clang::Decl::getDeclKindName()) != "TranslationUnit")
        {
            return true;
        }

        NotRelevantStmtVisiter SV(context, manager, resource, FD, &lockList, globalVarDeclSet, lockFindList, unlockFindList);
        SV.HandleTranslationUnit(*context);
        return true;
    }
};

void LockGranularityChecker::check()
{
    readConfig();
    //getEntryFunc();
    analyse();
}
void LockGranularityChecker::readConfig()
{
    std::unordered_map<std::string, std::string> ptrConfig =
        configure->getOptionBlock("ReduceLockChecker");
    vector<string> lock = parse(ptrConfig.find("lock")->second);
    for (auto it : lock)
    {
        lockFindList.insert(it);
    }
    vector<string> unlock = parse(ptrConfig.find("unlock")->second);
    for (auto it : unlock)
    {
        unlockFindList.insert(it);
    }
}

vector<string> LockGranularityChecker::parse(string funcs)
{
    vector<string> tmp;
    int param;
    int pos_begin = 0;
    int pos_end = 0;
    int size = funcs.size();
    string func_param;
    pos_end = funcs.find(":");
    while (pos_end != string::npos)
    {
        func_param = funcs.substr(0, pos_end);
        funcs = funcs.substr(pos_end + 1);
        pos_end = funcs.find(":");
        tmp.push_back(func_param);
    }
    return tmp;
}

void LockGranularityChecker::getEntryFunc()
{
    std::vector<ASTFunction *> topLevelFuncs = call_graph->getTopLevelFunctions();
    for (auto fun : topLevelFuncs)
    {
        const FunctionDecl *funDecl = manager->getFunctionDecl(fun);
        if (funDecl->getQualifiedNameAsString() == "main")
        {
            entryFunc = fun;
            return;
        }
    }
    entryFunc = nullptr;
    return;
}

void LockGranularityChecker::analyse()
{
    vector<ASTFile *> astFiles = resource->getASTFiles();
    vector<LockInfo *> report;
    std::set<VarDecl *> globalVarDeclSet;
    int size = astFiles.size();
    int read = 0;
    int count = 1;
    for (ASTFile *astFile : astFiles)
    {
        process_bar(float(count++) / astFiles.size() / 2);
        read++;
        ASTUnit *astUnit = manager->getASTUnit(astFile);
        ASTContext &AST = astUnit->getASTContext();
        GlobalVarDecl globVar;
        globVar.HandleTranslationUnit(AST);
        std::set<VarDecl *> &s2 = globVar.getSet();
        globalVarDeclSet.insert(s2.begin(), s2.end());
    }
    read = 0;
    for (ASTFile *astFile : astFiles)
    {
        process_bar(float(count++) / astFiles.size() / 2);
        read++;

        ASTUnit *astUnit = manager->getASTUnit(astFile);
        ASTContext &AST = astUnit->getASTContext();
        LockVisiter visiter(&AST, manager, resource, globalVarDeclSet, lockFindList, unlockFindList);
        visiter.HandleTranslationUnit(AST);
        vector<LockInfo *> &info = visiter.getList();
        report.insert(report.end(), info.begin(), info.end());
    }
    int i = 1;

    for (auto info : report)
    {

        LockInfo it = *info;
        if (it.report.size() == 0)
        {
            continue;
        }

        for (auto iter : it.report)
        {
            string desc = "";
            string fileName = iter.fileName;
            unsigned line = iter.line;
            unsigned col = iter.col;
            string var = iter.var;
            desc += "Lock from " + to_string(it.beginLine) + " to " + to_string(it.endLine);

            ReportSpace::Defect *def = new ReportSpace::Defect(ReportSpace::DefectName::ReduceLock,
                                                               ReportSpace::DefectType::Warning,
                                                               fileName, line, col,
                                                               "");
            def->addDesc(desc);
            external_bug_report.addToReport(def);
        }
        delete (info);
    }
}